// Udacity HW 6
// Poisson Blending Reference Calculation
// 本文件实现了泊松图像融合的串行参考版本（CPU）
// 用于将源图像中的非白色区域平滑地融合到目标图像中

#include "utils.h"           // 包含工具函数，如 make_uint2
#include <thrust/host_vector.h> // 虽未使用，但可能用于未来 GPU 版本的兼容

// 用于单次迭代求解泊松方程 Δf = g
// 使用 Jacobi 迭代法，边界条件来自目标图像
void computeIteration(
    const unsigned char* const dstImg,            // 目标图像像素值（用于边界条件）
    const unsigned char* const strictInteriorPixels, // 标记是否为“严格内部像素”(1=是)
    const unsigned char* const borderPixels,      // 标记是否为“边界像素”(1=是)
    const std::vector<uint2>& interiorPixelList,  // 所有内部像素坐标列表
    const size_t numColsSource,                   // 源图像宽度（列数）
    const float* const f,                         // 当前迭代步的解 f^{(k)}
    const float* const g,                         // 预计算的拉普拉斯项 ∇²s
    float* const f_next                           // 输出：下一步的解 f^{(k+1)}
) {
    // 【注意】这一行计算了第一个内部像素的偏移量，但后续未使用，可能是调试残留
    // unsigned int off = interiorPixelList[0].x * numColsSource + interiorPixelList[0].y;

    // 遍历每一个“严格内部像素”，进行泊松方程的一次迭代更新
    for (size_t i = 0; i < interiorPixelList.size(); ++i) {
        float blendedSum = 0.f;  // 存储来自其他“内部像素”的 f 值之和
        float borderSum  = 0.f;  // 存储来自“边界或外部像素”的 dstImg 值之和（作为边界条件）

        uint2 coord = interiorPixelList[i];  // 获取当前内部像素的 (行, 列) 坐标
        unsigned int offset = coord.x * numColsSource + coord.y; // 转换为一维索引

        // 检查左邻居 (offset - 1)
        if (strictInteriorPixels[offset - 1]) {
            // 如果左邻居也是内部像素，则使用当前解 f[offset-1]
            blendedSum += f[offset - 1];
        } else {
            // 否则它是边界或外部像素，使用目标图像的像素值作为狄利克雷边界条件
            borderSum += dstImg[offset - 1];
        }

        // 检查右邻居 (offset + 1)
        if (strictInteriorPixels[offset + 1]) {
            blendedSum += f[offset + 1];
        } else {
            borderSum += dstImg[offset + 1];
        }

        // 检查上邻居 (offset - numColsSource)
        if (strictInteriorPixels[offset - numColsSource]) {
            blendedSum += f[offset - numColsSource];
        } else {
            borderSum += dstImg[offset - numColsSource];
        }

        // 检查下邻居 (offset + numColsSource)
        if (strictInteriorPixels[offset + numColsSource]) {
            blendedSum += f[offset + numColsSource];
        } else {
            borderSum += dstImg[offset + numColsSource];
        }

        // 计算泊松方程的离散解：
        // f_next = (sum_内部邻居 f + sum_边界邻居 D + g) / 4
        // 其中 g = 4*s - s上 - s下 - s左 - s右 = ∇²s
        float f_next_val = (blendedSum + borderSum + g[offset]) / 4.f;

        // 将结果裁剪到 [0, 255] 范围内，确保是合法的像素值
        f_next[offset] = std::min(255.f, std::max(0.f, f_next_val));
    }
}

// 预计算泊松方程右边项 g = ∇²source（源图像的拉普拉斯）
// 该值在迭代过程中不变，只需计算一次
void computeG(
    const unsigned char* const channel,           // 某颜色通道的源图像数据
    float* const g,                               // 输出：g[i] = ∇²source[i]
    const size_t numColsSource,                   // 源图像宽度
    const std::vector<uint2>& interiorPixelList   // 内部像素列表
) {
    // 遍历所有严格内部像素
    for (size_t i = 0; i < interiorPixelList.size(); ++i) {
        uint2 coord = interiorPixelList[i];      // 当前像素坐标
        unsigned int offset = coord.x * numColsSource + coord.y; // 一维索引

        // 初始化为 4 倍中心像素值（拉普拉斯核中心系数）
        float sum = 4.f * channel[offset];

        // 减去四个方向邻居的像素值（标准 5 点拉普拉斯模板）
        sum -= (float)channel[offset - 1];       // 左
        sum -= (float)channel[offset + 1];       // 右
        sum -= (float)channel[offset - numColsSource]; // 上
        sum -= (float)channel[offset + numColsSource]; // 下

        // 存储计算结果到 g 数组
        g[offset] = sum;
    }
}

// 主函数：执行泊松图像融合
void reference_calc(
    const uchar4* const h_sourceImg,             // 源图像（RGBA）
    const size_t numRowsSource,                  // 源图像行数（高度）
    const size_t numColsSource,                  // 源图像列数（宽度）
    const uchar4* const h_destImg,               // 目标图像（背景图）
    uchar4* const h_blendedImg                   // 输出：融合后的图像
) {
    // 计算源图像总像素数
    size_t srcSize = numRowsSource * numColsSource;

    // 分配内存用于存储“前景掩码”
    // mask[i] = 1 表示该像素属于要融合的前景区域（非纯白）
    unsigned char* mask = new unsigned char[srcSize];

    // 构建掩码：判断源图像中哪些像素不是纯白色（R=G=B=255）
    for (int i = 0; i < srcSize; ++i) {
        // 如果 RGB 三通道之和小于 3*255=765，说明不是纯白
        mask[i] = (h_sourceImg[i].x + h_sourceImg[i].y + h_sourceImg[i].z < 3 * 255) ? 1 : 0;
    }

    // 分配内存用于标记“边界像素”和“严格内部像素”
    unsigned char *borderPixels = new unsigned char[srcSize];         // 边界像素标记
    unsigned char *strictInteriorPixels = new unsigned char[srcSize]; // 严格内部像素标记

    // 存储所有严格内部像素坐标的列表，用于加速迭代
    std::vector<uint2> interiorPixelList;

    // 遍历图像内部区域（排除边缘，防止越界）
    for (size_t r = 1; r < numRowsSource - 1; ++r) {
        for (size_t c = 1; c < numColsSource - 1; ++c) {
            unsigned int offset = r * numColsSource + c; // 当前像素一维索引

            // 如果当前像素在前景中
            if (mask[offset]) {
                // 检查四个邻居是否都在前景中
                bool isInterior =
                    mask[(r - 1) * numColsSource + c] &&  // 上
                    mask[(r + 1) * numColsSource + c] &&  // 下
                    mask[r * numColsSource + c - 1] &&    // 左
                    mask[r * numColsSource + c + 1];      // 右

                if (isInterior) {
                    // 是严格内部像素
                    strictInteriorPixels[offset] = 1;
                    borderPixels[offset] = 0;
                    interiorPixelList.push_back(make_uint2(r, c)); // 记录坐标
                } else {
                    // 是边界像素（在前景中但邻接背景）
                    strictInteriorPixels[offset] = 0;
                    borderPixels[offset] = 1;
                }
            } else {
                // 不在前景中，既不是内部也不是边界
                strictInteriorPixels[offset] = 0;
                borderPixels[offset] = 0;
            }
        }
    }

    // 分离源图像的 R、G、B 通道
    unsigned char* red_src   = new unsigned char[srcSize];
    unsigned char* green_src = new unsigned char[srcSize];
    unsigned char* blue_src  = new unsigned char[srcSize];

    for (int i = 0; i < srcSize; ++i) {
        red_src[i]   = h_sourceImg[i].x;   // R 通道
        green_src[i] = h_sourceImg[i].y;   // G 通道
        blue_src[i]  = h_sourceImg[i].z;   // B 通道
    }

    // 分离目标图像的 R、G、B 通道
    unsigned char* red_dst   = new unsigned char[srcSize];
    unsigned char* green_dst = new unsigned char[srcSize];
    unsigned char* blue_dst  = new unsigned char[srcSize];

    for (int i = 0; i < srcSize; ++i) {
        red_dst[i]   = h_destImg[i].x;
        green_dst[i] = h_destImg[i].y;
        blue_dst[i]  = h_destImg[i].z;
    }

    // 分配内存用于存储预计算的 g = ∇²source
    float *g_red   = new float[srcSize];
    float *g_green = new float[srcSize];
    float *g_blue  = new float[srcSize];

    // 初始化 g 数组为 0
    memset(g_red,   0, srcSize * sizeof(float));
    memset(g_green, 0, srcSize * sizeof(float));
    memset(g_blue,  0, srcSize * sizeof(float));

    // 预计算每个通道的拉普拉斯值
    computeG(red_src,   g_red,   numColsSource, interiorPixelList);
    computeG(green_src, g_green, numColsSource, interiorPixelList);
    computeG(blue_src,  g_blue,  numColsSource, interiorPixelList);

    // 为每个颜色通道分配两个缓冲区，用于 Jacobi 迭代（双缓冲）
    float *blendedValsRed_1   = new float[srcSize];
    float *blendedValsRed_2   = new float[srcSize];
    float *blendedValsGreen_1 = new float[srcSize];
    float *blendedValsGreen_2 = new float[srcSize];
    float *blendedValsBlue_1  = new float[srcSize];
    float *blendedValsBlue_2  = new float[srcSize];

    // 初始化迭代解为源图像像素值（初始猜测）
    for (size_t i = 0; i < srcSize; ++i) {
        blendedValsRed_1[i]   = red_src[i];
        blendedValsRed_2[i]   = red_src[i];
        blendedValsGreen_1[i] = green_src[i];
        blendedValsGreen_2[i] = green_src[i];
        blendedValsBlue_1[i]  = blue_src[i];
        blendedValsBlue_2[i]  = blue_src[i];
    }

    // 设置迭代次数（经验性选择，通常 800 次可收敛）
    const size_t numIterations = 800;

    // === 红色通道求解 ===
    for (size_t i = 0; i < numIterations; ++i) {
        // 使用 blendedValsRed_1 作为输入，计算 blendedValsRed_2
        computeIteration(
            red_dst, strictInteriorPixels, borderPixels,
            interiorPixelList, numColsSource,
            blendedValsRed_1, g_red, blendedValsRed_2
        );
        // 交换指针，使下一轮使用新值作为输入
        std::swap(blendedValsRed_1, blendedValsRed_2);
    }

    // === 绿色通道求解 ===
    for (size_t i = 0; i < numIterations; ++i) {
        computeIteration(
            green_dst, strictInteriorPixels, borderPixels,
            interiorPixelList, numColsSource,
            blendedValsGreen_1, g_green, blendedValsGreen_2
        );
        std::swap(blendedValsGreen_1, blendedValsGreen_2);
    }

    // === 蓝色通道求解 ===
    for (size_t i = 0; i < numIterations; ++i) {
        computeIteration(
            blue_dst, strictInteriorPixels, borderPixels,
            interiorPixelList, numColsSource,
            blendedValsBlue_1, g_blue, blendedValsBlue_2
        );
        std::swap(blendedValsBlue_1, blendedValsBlue_2);
    }

    // 经过最后一次 swap 后，最终结果其实在 *_1 中
    // 所以再 swap 一次，确保最终结果在 *_2 中（与下面赋值逻辑一致）
    std::swap(blendedValsRed_1,   blendedValsRed_2);
    std::swap(blendedValsGreen_1, blendedValsGreen_2);
    std::swap(blendedValsBlue_1,  blendedValsBlue_2);

    // 先将整个目标图像复制到输出图像（背景保持不变）
    memcpy(h_blendedImg, h_destImg, sizeof(uchar4) * srcSize);

    // 将融合区域（内部像素）的结果写入输出图像
    for (size_t i = 0; i < interiorPixelList.size(); ++i) {
        uint2 coord = interiorPixelList[i];
        unsigned int offset = coord.x * numColsSource + coord.y;

        // 写入融合后的 RGB 值（已裁剪到 [0,255]）
        h_blendedImg[offset].x = static_cast<unsigned char>(blendedValsRed_2[offset]);
        h_blendedImg[offset].y = static_cast<unsigned char>(blendedValsGreen_2[offset]);
        h_blendedImg[offset].z = static_cast<unsigned char>(blendedValsBlue_2[offset]);
        // .w (alpha) 保持不变，来自 h_destImg
    }

    // 释放所有动态分配的内存，防止内存泄漏
    delete[] mask;
    delete[] red_src;     delete[] red_dst;
    delete[] green_src;   delete[] green_dst;
    delete[] blue_src;    delete[] blue_dst;
    delete[] g_red;       delete[] g_green;     delete[] g_blue;
    delete[] blendedValsRed_1;   delete[] blendedValsRed_2;
    delete[] blendedValsGreen_1; delete[] blendedValsGreen_2;
    delete[] blendedValsBlue_1;  delete[] blendedValsBlue_2;
    delete[] borderPixels;
    delete[] strictInteriorPixels;

    // 函数结束，输出图像 h_blendedImg 已被修改
}